from urllib.parse import parse_qs
from channels.exceptions import DenyConnection
from channels.middleware import BaseMiddleware


class JWTAuthMiddleware(BaseMiddleware):
    async def __call__(self, scope, receive, send):
        from HomeAutoAI.common.utils.Utils import async_get_user_from_token

        # 从路径查询参数中获取 token
        query_string = scope.get('query_string', b'').decode()  # 获取查询字符串
        params = parse_qs(query_string)  # 解析查询字符串
        token = params.get('token', [None])[0]  # 获取 `token` 参数的值

        print(f"WebSocket query string: {query_string}")
        print(f"Extracted token: {token}")

        # 如果 token 存在，验证并获取用户
        if token:
            user = await async_get_user_from_token(token)
            if user:
                scope['user'] = user  # 将用户信息注入到 scope 中
            else:
                raise DenyConnection("未认证")
        else:
            raise DenyConnection("未提供Token")

        # 继续处理连接
        return await super().__call__(scope, receive, send)
